
package com.ship.core.processor;

import com.ship.core.annotation.RequestDecryptBody;
import com.ship.core.annotation.ResponseEncryptBody;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.Conventions;
import org.springframework.core.MethodParameter;
import org.springframework.core.annotation.AnnotatedElementUtils;
import org.springframework.http.HttpInputMessage;
import org.springframework.http.MediaType;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.converter.HttpMessageNotReadableException;
import org.springframework.http.converter.HttpMessageNotWritableException;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.http.server.ServletServerHttpResponse;
import org.springframework.util.Assert;
import org.springframework.validation.BindingResult;
import org.springframework.web.HttpMediaTypeNotAcceptableException;
import org.springframework.web.HttpMediaTypeNotSupportedException;
import org.springframework.web.accept.ContentNegotiationManager;
import org.springframework.web.accept.PathExtensionContentNegotiationStrategy;
import org.springframework.web.bind.MethodArgumentNotValidException;
import org.springframework.web.bind.WebDataBinder;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.ResponseBody;
import org.springframework.web.bind.support.WebDataBinderFactory;
import org.springframework.web.context.request.NativeWebRequest;
import org.springframework.web.method.support.ModelAndViewContainer;
import org.springframework.web.servlet.mvc.method.annotation.AbstractMessageConverterMethodProcessor;

import javax.servlet.http.HttpServletRequest;
import java.io.IOException;
import java.lang.reflect.Type;
import java.util.*;

/**
 * 自定义注解处理器
 */
public class RequestResponseBodyMethodProcessor extends AbstractMessageConverterMethodProcessor {

    private final Logger logger = LoggerFactory.getLogger(getClass());

    private MyRequestResponseBodyAdviceChain advice;

    private List<MediaType> allSupportedMediaTypes;
    private List<HttpMessageConverter<?>> messageConverters;
    private ContentNegotiationManager contentNegotiationManager;


    private final Set<String> safeExtensions = new HashSet<String>();


    private static final Set<String> WHITELISTED_EXTENSIONS = new HashSet<>(Arrays.asList(
            "txt", "text", "yml", "properties", "csv",
            "json", "xml", "atom", "rss",
            "png", "jpe", "jpeg", "jpg", "gif", "wbmp", "bmp"));

    private PathExtensionContentNegotiationStrategy pathStrategy;

    public RequestResponseBodyMethodProcessor(List<HttpMessageConverter<?>> converters) {
        super(converters);
        this.messageConverters = converters;
    }

    public RequestResponseBodyMethodProcessor(List<HttpMessageConverter<?>> converters, ContentNegotiationManager manager, List<Object> requestResponseBodyAdvice) {
        super(converters, manager, requestResponseBodyAdvice);
        Assert.notEmpty(converters, "'messageConverters' must not be empty");
        this.messageConverters = converters;
        this.allSupportedMediaTypes = getAllSupportedMediaTypes(converters);
        this.advice = new MyRequestResponseBodyAdviceChain(requestResponseBodyAdvice);
        this.contentNegotiationManager = (manager != null ? manager : new ContentNegotiationManager());
        this.pathStrategy = initPathStrategy(this.contentNegotiationManager);
        this.safeExtensions.addAll(this.contentNegotiationManager.getAllFileExtensions());
        this.safeExtensions.addAll(WHITELISTED_EXTENSIONS);
    }


    @Override
    public boolean supportsParameter(MethodParameter parameter) {
        boolean b = parameter.hasParameterAnnotation(RequestDecryptBody.class);
        if (b) {
            return b;
        } else {
            return parameter.hasParameterAnnotation(RequestBody.class);
        }
    }

    @Override
    public Object resolveArgument(MethodParameter parameter, ModelAndViewContainer mavContainer,
                                  NativeWebRequest webRequest, WebDataBinderFactory binderFactory) throws Exception {
        Object arg = readWithMessageConverters(webRequest, parameter, parameter.getGenericParameterType());
        String name = Conventions.getVariableNameForParameter(parameter);

        WebDataBinder binder = binderFactory.createBinder(webRequest, arg, name);
        if (arg != null) {
            validateIfApplicable(binder, parameter);
            if (binder.getBindingResult().hasErrors() && isBindExceptionRequired(binder, parameter)) {
                throw new MethodArgumentNotValidException(parameter, binder.getBindingResult());
            }
        }
        mavContainer.addAttribute(BindingResult.MODEL_KEY_PREFIX + name, binder.getBindingResult());

        return arg;
    }

    @Override
    public boolean supportsReturnType(MethodParameter returnType) {

        boolean b = (AnnotatedElementUtils.hasAnnotation(returnType.getContainingClass(), ResponseEncryptBody.class) ||
                returnType.hasMethodAnnotation(ResponseEncryptBody.class));
        if (b) {
            return b;
        } else {
            return (AnnotatedElementUtils.hasAnnotation(returnType.getContainingClass(), ResponseBody.class) ||
                    returnType.hasMethodAnnotation(ResponseBody.class));
        }
    }

    @Override
    public void handleReturnValue(Object returnValue, MethodParameter returnType, ModelAndViewContainer mavContainer, NativeWebRequest webRequest) throws IOException, HttpMediaTypeNotAcceptableException, HttpMessageNotWritableException {
        mavContainer.setRequestHandled(true);
        ServletServerHttpRequest inputMessage = createInputMessage(webRequest);
        ServletServerHttpResponse outputMessage = createOutputMessage(webRequest);
        if (returnType.hasMethodAnnotation(ResponseEncryptBody.class)) {

            super.writeWithMessageConverters(returnValue, returnType, inputMessage, outputMessage);
        } else if (returnType.hasMethodAnnotation(ResponseBody.class)) {
            this.writeWithMessageConverters(returnValue, returnType, inputMessage, outputMessage);
        }
    }

    @Override
    protected <T> Object readWithMessageConverters(NativeWebRequest webRequest, MethodParameter methodParam, Type paramType) throws IOException, HttpMediaTypeNotSupportedException, HttpMessageNotReadableException {
        boolean b = methodParam.hasParameterAnnotation(RequestDecryptBody.class);
        HttpServletRequest servletRequest = webRequest.getNativeRequest(HttpServletRequest.class);
        ServletServerHttpRequest inputMessage = new ServletServerHttpRequest(servletRequest);
        if (b) {
            Object arg = super.readWithMessageConverters(inputMessage, methodParam, paramType);
            return arg;
        } else {
            Object arg = this.readWithMessageConverters(inputMessage, methodParam, paramType);
            return arg;
        }

    }

    @Override
    protected <T> Object readWithMessageConverters(HttpInputMessage inputMessage, MethodParameter param, Type targetType) throws IOException, HttpMediaTypeNotSupportedException, HttpMessageNotReadableException {
        RequestBodyRead requestBodyRead = new RequestBodyRead(messageConverters, advice, allSupportedMediaTypes);
        return requestBodyRead.readWithMessageConverters(inputMessage, param, targetType);
    }

    @Override
    protected <T> void writeWithMessageConverters(T value, MethodParameter returnType, ServletServerHttpRequest inputMessage, ServletServerHttpResponse outputMessage) throws IOException, HttpMediaTypeNotAcceptableException, HttpMessageNotWritableException {

        ResponseBodyWrite write = new ResponseBodyWrite(allSupportedMediaTypes, messageConverters, advice, contentNegotiationManager, safeExtensions, pathStrategy);
        write.writeWithMessageConverters(value, returnType, inputMessage, outputMessage);
    }

    private static List<MediaType> getAllSupportedMediaTypes(List<HttpMessageConverter<?>> messageConverters) {
        Set<MediaType> allSupportedMediaTypes = new LinkedHashSet<MediaType>();
        for (HttpMessageConverter<?> messageConverter : messageConverters) {
            allSupportedMediaTypes.addAll(messageConverter.getSupportedMediaTypes());
        }
        List<MediaType> result = new ArrayList<MediaType>(allSupportedMediaTypes);
        MediaType.sortBySpecificity(result);
        return Collections.unmodifiableList(result);
    }


    private static PathExtensionContentNegotiationStrategy initPathStrategy(ContentNegotiationManager manager) {
        Class<PathExtensionContentNegotiationStrategy> clazz = PathExtensionContentNegotiationStrategy.class;
        PathExtensionContentNegotiationStrategy strategy = manager.getStrategy(clazz);
        return (strategy != null ? strategy : new PathExtensionContentNegotiationStrategy());
    }

}
